import heapq
class Solution(object):
    def trapRainWater(self, heightMap):
        """
        :type heightMap: List[List[int]]
        :rtype: int
        """
        m = len(heightMap)
        n = len(heightMap[0])
        if m <= 2 or n <= 2:
            return 0
        used = [[0 for _ in range(n)] for _ in range(m)]
        stack = []
        for j in range(n):
            used[0][j] = 1
            stack.append([heightMap[0][j], 0, j])
            used[m - 1][j] = 1
            stack.append([heightMap[m - 1][j], m - 1, j])
        for i in range(1, m - 1):
            used[i][0] = 1
            stack.append([heightMap[i][0], i, 0])
            used[i][n - 1] = 1
            stack.append([heightMap[i][n - 1], i, n - 1])
        count = (m - 2) * (n - 2)
        heapq.heapify(stack)
        total = 0
        while count > 0:
            node = heapq.heappop(stack)
            x = node[1]
            y = node[2]
            for ux, uy in [(x - 1, y), (x + 1,y), (x, y - 1), (x, y + 1)]:
                if 0 <= ux < m and 0 <= uy < n and not used[ux][uy]:
                    total += max(node[0] - heightMap[ux][uy], 0)
                    heightMap[ux][uy] = max(heightMap[ux][uy], node[0])
                    used[ux][uy] = 1
                    heapq.heappush(stack, [heightMap[ux][uy], ux, uy])
                    count -= 1
        return total


data = Solution()
heightMap = [[1, 4, 3, 1, 3, 2], [3, 2, 1, 3, 2, 4], [2, 3, 3, 2, 3, 1]]
print(data.trapRainWater(heightMap))
heightMap = [[3, 3, 3, 3, 3], [2, 2, 2, 2, 2], [3, 2, 1, 2, 3], [3, 2, 2, 2, 3], [3, 3, 3, 3, 3]]
print(data.trapRainWater(heightMap))
heightMap = [[12, 13, 1, 12], [13, 4, 13, 12], [13, 8, 10, 12], [12, 13, 12, 12], [13, 13, 13, 13]]
print(data.trapRainWater(heightMap))
heightMap = [[14, 20, 11, 19, 19, 16], [11, 10, 7, 4, 9, 6], [17, 2, 2, 6, 10, 9], [15, 9, 2, 1, 4, 1],
             [15, 5, 5, 5, 8, 7], [14, 2, 8, 6, 10, 7]]
print(data.trapRainWater(heightMap))
